#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2024-08-06


"""
【节点名称】：
    MovingObjectDetNode
【依赖项安装】：
    pip install spirems
    pip install ultralytics
【订阅类型】：
    sensor_msgs::CompressedImage （输入图像）
【发布类型】：
    spirecv_msgs::2DTargets （检测结果）
    sensor_msgs::CompressedImage （可视化结果，需借助可视化工具）
    std_msgs::Boolean （如果输入节点是数据集，则检测完成发布该话题让输入节点继续工作）
【构造参数说明】：
    parameter_file (str): 全局参数文件
    sms_shutdown (bool): 是否接收全局关闭信号，如果需要长期后台执行，建议设置为False
    specified_input_topic (str): 指定输入的话题地址
    specified_output_topic (str): 指定输出的话题地址
    realtime_det (bool): 是否是实时检测器，设置为True可以降低延迟，但可能会产生丢帧
【节点参数】：
    confidence (float): 目标得分阈值
    nms_thresh (float): NMS后处理参数
    dataset_name (str): 数据集名称
    pt_model (str): 加载模型名称
【备注】：
    无
"""
import cv2
from ultralytics import YOLO
import threading
from queue import Queue
import numpy as np
import time
from typing import Union
from spirems import Subscriber, Publisher, cvimg2sms, sms2cvimg, def_msg, QoS
from spirecv.base.BaseNode import BaseNode
from spirecv.algorithm.moving_object_det.GLMD.MOD2 import MOD2_global


def trans_det_results(det_results):
    sms_results = def_msg('spirecv_msgs::2DTargets')

    sms_results["file_name"] = ""
    sms_results["height"] = int(det_results.orig_shape[0])
    sms_results["width"] = int(det_results.orig_shape[1])
    sms_results["targets"] = []
    cls = det_results.boxes.cls.cpu().numpy().astype(np.float64)
    conf = det_results.boxes.conf.cpu().numpy().astype(np.float64)
    xywh = det_results.boxes.xywh.cpu().numpy().astype(np.float64)
    for i in range(len(cls)):
        ann = dict()
        ann["category_name"] = det_results.names[int(cls[i])].strip().replace(' ', '_').lower()
        ann["category_id"] = int(cls[i])
        ann["score"] = round(conf[i], 3)
        ann["bbox"] = [round(j, 3) for j in xywh[i].tolist()]
        ann["bbox"][0] = ann["bbox"][0] - ann["bbox"][2] / 2
        ann["bbox"][1] = ann["bbox"][1] - ann["bbox"][3] / 2
        sms_results["targets"].append(ann)

    return sms_results


class MovingObjectDetNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown
        )
        input_url = '/' + job_name + '/sensor/image_raw'
        output_url = '/' + job_name + '/detector/results'

        self.job_queue = Queue()
        self.queue_pool.append(self.job_queue)

        self._image_reader = Subscriber(
            input_url, 'sensor_msgs::CompressedImage', self.image_callback,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._result_writer = Publisher(
            output_url, 'spirecv_msgs::2DTargets',
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._show_writer = Publisher(
            '/' + job_name + '/detector/image_results', 'sensor_msgs::CompressedImage',
            ip=ip, port=port
        )

        self.last_img = None
        self.last_img_name = None
        self.last_vid_name = None

        self.infer_n_ims = 0
        self.infer_time3 = 0.0
        self.infer_time2 = 0.0
        self.infer_time1 = 0.0

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()
        self._result_writer.kill()
        self._show_writer.kill()

    def image_callback(self, msg):
        self.job_queue.put(msg)

    def run(self):
        while self.is_running():
            msg = self.job_queue.get(block=True)
            if msg is None:
                break

            t1 = time.time()
            file_name = msg['file_name'] if 'file_name' in msg else ''
            img = sms2cvimg(msg)

            assert 'video_name' in msg

            # BEGIN
            if self.last_img is not None and self.last_vid_name is not None and msg['video_name'] == self.last_vid_name:
                print("Frame1: {}, Frame2: {}".format(self.last_img_name, msg['file_name']))
                print("Shape1: {}, Shape2: {}".format(self.last_img.shape, img.shape))

                rect_final = MOD2_global(self.last_img, img)
                if len(rect_final):
                    cv2.rectangle(
                        img,
                        (int(rect_final[0]), int(rect_final[1])),
                        (int(rect_final[0] + rect_final[2]), int(rect_final[1] + rect_final[3])),
                        (0, 0, 255),
                        2,
                        cv2.LINE_AA
                    )
                cv2.imshow('img', img)
                cv2.waitKey(5)
                print(rect_final)
            # END

            self.last_img = img
            self.last_img_name = msg['file_name']
            self.last_vid_name = msg['video_name']

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parameter_file = '../../../params/spirecv2/default_params.json'
    node = MovingObjectDetNode('MOD', param_dict_or_file=parameter_file)
